{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# HybridGaussianFactor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/hybrid/doc/HybridGaussianFactor.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `HybridGaussianFactor` represents a factor where the underlying probability distribution is Gaussian, but *which* Gaussian distribution applies depends on the assignment of a set of discrete variables. It inherits from `HybridFactor`.\n",
        "\n",
        "Mathematically, it represents a factor $\\phi(X, M)$ over continuous variables $X$ and discrete variables $M$, defined as:\n",
        "$$\n",
        "\\phi(X, M=m) \\propto \\exp(-E_m - \\frac12 ||A_m X - b_m||^2)\n",
        "$$\n",
        "where for each discrete assignment $m$, we have a standard Gaussian factor defined by $A_m$, $b_m$ (and an associated noise model, implicitly defining the Mahalanobis distance), plus an optional scalar energy term $E_m$.\n",
        "\n",
        "Internally, it uses a `DecisionTree<Key, GaussianFactorValuePair>` to store the different Gaussian factor components (`GaussianFactor::shared_ptr`) and their associated scalar energies (`double`), indexed by the `DiscreteKey`s.\n",
        "\n",
        "These factors typically arise from:\n",
        "1.  Linearizing a `HybridNonlinearFactor`.\n",
        "2.  Representing sensor models with discrete operating modes (e.g., a sensor that works differently in rain vs. sun).\n",
        "3.  As intermediate results during hybrid elimination."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "\n",
        "from gtsam import HybridGaussianFactor, JacobianFactor, DiscreteKeys, DiscreteValues\n",
        "from gtsam.symbol_shorthand import X, D"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Can be initialized with one or more discrete keys and the corresponding Gaussian factor components."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HybridGaussianFactor (factors only):\n",
            "HybridGaussianFactor\n",
            "\n",
            "Hybrid [x0; d0]{\n",
            " Choice(d0) \n",
            " 0 Leaf :\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 0 ]\n",
            "  Noise model: unit (1) \n",
            "scalar: 0\n",
            "\n",
            " 1 Leaf :\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 5 ]\n",
            "isotropic dim=1 sigma=2\n",
            "scalar: 0\n",
            "\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "# --- Example 1: Single Discrete Key ---\n",
        "dk0 = (D(0), 2) # Binary key D0\n",
        "\n",
        "# Gaussian factor for mode D0=0: N(X(0); mean=0, sigma=1) -> ||x0 - 0||^2 / 2\n",
        "gf0 = JacobianFactor(X(0), np.eye(1), np.zeros(1), gtsam.noiseModel.Isotropic.Sigma(1, 1.0))\n",
        "# Gaussian factor for mode D0=1: N(X(0); mean=5, sigma=2) -> ||x0 - 5||^2 / (2*4)\n",
        "gf1 = JacobianFactor(X(0), np.eye(1), np.array([5.0]), gtsam.noiseModel.Isotropic.Sigma(1, 2.0))\n",
        "\n",
        "# Option A: Just factors (scalar energy E_m defaults to 0)\n",
        "hgf_a = HybridGaussianFactor(dk0, [gf0, gf1])\n",
        "print(\"HybridGaussianFactor (factors only):\")\n",
        "hgf_a.print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "65a3f460",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "HybridGaussianFactor (factors + scalars):\n",
            "HybridGaussianFactor\n",
            "\n",
            "Hybrid [x0; d0]{\n",
            " Choice(d0) \n",
            " 0 Leaf :\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 0 ]\n",
            "  Noise model: unit (1) \n",
            "scalar: 0.356675\n",
            "\n",
            " 1 Leaf :\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 5 ]\n",
            "isotropic dim=1 sigma=2\n",
            "scalar: 1.20397\n",
            "\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "# Option B: Factors with scalar energies E_m\n",
        "# Example: Add -log P(D0=0)= -log(0.7) and -log P(D0=1)=-log(0.3)\n",
        "scalar0 = -np.log(0.7)\n",
        "scalar1 = -np.log(0.3)\n",
        "hgf_b = HybridGaussianFactor(dk0, [(gf0, scalar0), (gf1, scalar1)])\n",
        "print(\"\\nHybridGaussianFactor (factors + scalars):\")\n",
        "hgf_b.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3a9d47a1",
      "metadata": {},
      "source": [
        "In C++ you can create HybridGaussianFactors with multiple discrete keys as well, but that capability is not yet exposed in the wrapper."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Accessing Components\n",
        "\n",
        "You can retrieve the underlying decision tree structure and select a specific Gaussian factor component based on a discrete assignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Factor/Scalar pair for D0=0:\n",
            "Factor:\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 0 ]\n",
            "  Noise model: unit (1) \n",
            "Scalar: 0.35667494393873245\n",
            "\n",
            "Factor/Scalar pair for D0=1:\n",
            "Factor:\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 5 ]\n",
            "isotropic dim=1 sigma=2\n",
            "Scalar: 1.2039728043259361\n"
          ]
        }
      ],
      "source": [
        "# Using hgf_b from Example 1B\n",
        "\n",
        "# Get the factor/scalar pair for a specific assignment\n",
        "assignment0 = gtsam.DiscreteValues()\n",
        "assignment0[D(0)] = 0\n",
        "factor_pair0 = hgf_b(assignment0) # Calls operator()\n",
        "print(\"\\nFactor/Scalar pair for D0=0:\")\n",
        "factor_pair0[0].print(\"Factor:\") # GaussianFactor component\n",
        "print(f\"Scalar: {factor_pair0[1]}\") # Scalar energy term\n",
        "\n",
        "assignment1 = gtsam.DiscreteValues()\n",
        "assignment1[D(0)] = 1\n",
        "factor_pair1 = hgf_b(assignment1)\n",
        "print(\"\\nFactor/Scalar pair for D0=1:\")\n",
        "factor_pair1[0].print(\"Factor:\")\n",
        "print(f\"Scalar: {factor_pair1[1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Error Calculation (`error`, `errorTree`)\n",
        "\n",
        "*   `error(HybridValues)`: Calculates the total error $E_m + \\frac{1}{2} ||A_m X - b_m||^2$ for the specified continuous values $X$ and discrete assignment $M=m$.\n",
        "*   `errorTree(VectorValues)`: Calculates the Gaussian error term $\\frac{1}{2} ||A_m X - b_m||^2$ for the given continuous values $X$ across *all* discrete assignments $m$, returning the results as an `AlgebraicDecisionTree` indexed by the discrete keys $M$. It also adds the scalar energy $E_m$ to each leaf."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Error for X0=0.1, D0=0: 0.36167494393873245 (Expected ~ 0.36167494393873245)\n",
            "Error for X0=4.0, D0=1: 1.3289728043259361 (Expected ~ 1.3289728043259361)\n",
            "\n",
            "Error Tree for X0=1.0:\n",
            "ErrorTree\n",
            " Choice(d0) \n",
            "ErrorTree\n",
            " 0 Leaf 0.85667494\n",
            "ErrorTree\n",
            " 1 Leaf 3.2039728\n",
            "(Expected Leaf 0 ~ 0.8566749439387324)\n",
            "(Expected Leaf 1 ~ 3.203972804325936)\n"
          ]
        }
      ],
      "source": [
        "# Using hgf_b = HybridGaussianFactor(dk0, [(gf0, scalar0), (gf1, scalar1)])\n",
        "# gf0: N(X(0); 0, 1), scalar0 = -log(0.7)\n",
        "# gf1: N(X(0); 5, 2), scalar1 = -log(0.3)\n",
        "\n",
        "# --- error(HybridValues) ---\n",
        "hybrid_vals0 = gtsam.HybridValues()\n",
        "hybrid_vals0.insert(X(0), np.array([0.1])) # Continuous value\n",
        "hybrid_vals0.insert(D(0), 0)              # Discrete assignment\n",
        "\n",
        "err0 = hgf_b.error(hybrid_vals0)\n",
        "# Expected: scalar0 + 0.5 * ||0.1 - 0||^2 / 1^2\n",
        "print(f\"\\nError for X0=0.1, D0=0: {err0} (Expected ~ {-np.log(0.7) + 0.5*0.1**2})\")\n",
        "\n",
        "hybrid_vals1 = gtsam.HybridValues()\n",
        "hybrid_vals1.insert(X(0), np.array([4.0])) # Continuous value\n",
        "hybrid_vals1.insert(D(0), 1)              # Discrete assignment\n",
        "\n",
        "err1 = hgf_b.error(hybrid_vals1)\n",
        "# Expected: scalar1 + 0.5 * ||4.0 - 5.0||^2 / 2^2\n",
        "print(f\"Error for X0=4.0, D0=1: {err1} (Expected ~ {-np.log(0.3) + 0.5*(-1.0)**2 / 4.0})\")\n",
        "\n",
        "# --- errorTree(VectorValues) ---\n",
        "cont_vals = gtsam.VectorValues()\n",
        "cont_vals.insert(X(0), np.array([1.0])) # Evaluate Gaussian part at X0=1.0\n",
        "\n",
        "adt = hgf_b.errorTree(cont_vals)\n",
        "# Leaf 0 (D0=0): scalar0 + 0.5 * ||1.0 - 0||^2 / 1^2 = scalar0 + 0.5\n",
        "# Leaf 1 (D0=1): scalar1 + 0.5 * ||1.0 - 5.0||^2 / 2^2 = scalar1 + 0.5 * 16 / 4 = scalar1 + 2.0\n",
        "print(\"\\nError Tree for X0=1.0:\")\n",
        "adt.print()\n",
        "print(f\"(Expected Leaf 0 ~ {-np.log(0.7) + 0.5})\")\n",
        "print(f\"(Expected Leaf 1 ~ {-np.log(0.3) + 2.0})\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Restriction (`restrict`)\n",
        "\n",
        "Fixes the discrete variables according to a given assignment, returning the corresponding `GaussianFactor` component (as a `gtsam.Factor::shared_ptr`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Restricted Factor for D0=1:\n",
            "Continuous [x0]{\n",
            " Leaf :\n",
            "  A[x0] = [\n",
            "\t1\n",
            "]\n",
            "  b = [ 5 ]\n",
            "isotropic dim=1 sigma=2\n",
            "scalar: 1.20397\n",
            "\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "# Using hgf_b\n",
        "restricted_factor_ptr = hgf_b.restrict(assignment1)# Choose mode 1\n",
        "\n",
        "if restricted_factor_ptr:\n",
        "    restricted_factor_ptr.print(\"\\nRestricted Factor for D0=1:\")\n",
        "    # Should be equivalent to gf1: N(X(0); 5, 2)\n",
        "else:\n",
        "    print(\"\\nRestriction failed.\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
